# 构建sft数据集
from enum import Enum
import json
 
class AnswerType(Enum):
    ALL = 1
    ALL_NOT = 2
    NORMAL = 3

with open('/home/llm_user/index/medical/llm-retrieval-qa/data/CMB-Exam/CMB-train/CMB-train-merge.json','r',encoding='utf-8') as f:
    train = json.load(f)
# with open('/home/llm_user/index/medical/llm-retrieval-qa/data/CMB-Exam/CMB-train/CMB-train-hierarchy.json','r',encoding='utf-8') as f:
#     train_hierarchy = json.load(f)
# 找中医相关的数据

def answer_type(answer):
    if '以上' in answer and ('不' in answer or '非' in answer) and len(answer) <= 5:
        return AnswerType.ALL_NOT
    elif '以上' in answer and len(answer) <= 5:
        return AnswerType.ALL
    else:
        return AnswerType.NORMAL

def process_question(question):
    question = question.replace('（）','')
    question = question.replace('()','')
    if '。' == question[-1]: question = question[:-1]
    return question

raw_data = []
sft_data = []
for item in train:
    if item['exam_class'] == '中医学与中药学':
        option_answers = item['answer']
        options = item['option']
        answers = []
        if len(option_answers) > 5:
            continue
    
        at = answer_type(options[option_answers[0]])
        if at == AnswerType.ALL:
            choosen = option_answers[0]
            option_answers = ['A','B','C','D','E']
            option_answers.remove(choosen)
        elif at == AnswerType.ALL_NOT:
            continue

        for _a in option_answers:
            answers.append(options[_a])
        answer = ';'.join(answers)

        sft_data.append(json.dumps({
            'instruction': process_question(item['question']),
            'input': '',
            'output':answer
        },ensure_ascii=False))
        raw_data.append(
            process_question(item['question']) + answer
        )


with open('/home/llm_user/index/medical/llm-retrieval-qa/data/CMB-Exam/CMB-train/CMB-train-sft.json','w',encoding='utf-8') as f:
    f.write(
        '\n'.join(sft_data)
    )

with open('/home/llm_user/index/medical/llm-retrieval-qa/data/text/CMB-train.txt','w',encoding='utf-8') as f:
    f.write(
        '\n'.join(raw_data)
    )


